import json
from zhipuai import ZhipuAI
from tool.prompt import RELATED_PROMPT
from sklearn.cluster import DBSCAN
import re
from FlagEmbedding import BGEM3FlagModel
from tqdm import tqdm
from generate_zhipu import generate_by_user_contents

def process_paper(paper: str):
    chunk = paper.split("#")
    if len(chunk) >= 3:
        return "#" + chunk[1] + "#" + chunk[2]
    else:
        return paper

def get_related(data, topic: str):
    from dotenv import load_dotenv
    load_dotenv()
    client = ZhipuAI()
    abs_all = ""
    user_contents = []
    for i in range(len(data)):
        if (i+1)%5 != 0:
            if len(data[i]) >= 1:
                if data[i][0].get('chunk_text') is not None:
                    abs_all = abs_all + "## id: " + str(i + 1) + "\n"
                    abs_all = abs_all + process_paper(data[i][0]["chunk_text"]) + "\n"

        else:
            if len(data[i]) >= 1:
                if data[i][0].get('chunk_text') is not None:
                    abs_all = abs_all + "## id: " + str(i + 1) + "\n"
                    abs_all = abs_all + process_paper(data[i][0]["chunk_text"]) + "\n"
            content = RELATED_PROMPT.format(topic=topic, abstract=abs_all)
            user_contents.append(content)
            abs_all = ""
    responses = generate_by_user_contents(client, user_contents, False, 48)
    related = []
    for response in responses:
        match = re.search(r'```json([\s\S]*?)```', response)

        if match:
            inner_content = match.group(1)
            related.extend(json.loads(inner_content))
        else:
            raise ValueError(
                "读取失败"
            )
    return related

if __name__ == "__main__":
    model = BGEM3FlagModel('BAAI/bge-m3',  use_fp16=True, cache_dir="./bge")
    #model = BGEM3FlagModel("/data/huggingface/models/BAAI/bge-m3/", use_fp16=True)
    with open("./files/loss_function_paper.json", "r")as f:
        data = json.load(f)
    related =  get_related(data, topic="损失函数")
    with open("related.json", "w")as f:
        json.dump(related, f, ensure_ascii=False, indent=4)